import io
from collections import defaultdict
from typing import List, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image

from util.box_ops import box_cxcywh_to_xyxy
from util.misc import  interpolate
import numpy as np
# try:
#     from panopticapi.utils import id2rgb, rgb2id
# except ImportError:
#     pass


class PostProcess(nn.Module):
    """ This module converts the model's output into the format expected by the coco api"""
    @torch.no_grad()
    def forward(self, outputs, target_sizes, type='detr'):
        out_logits, out_bbox = outputs['pred_logits'], outputs['pred_boxes']

        assert len(out_logits) == len(target_sizes)

        bs, query_len, category = out_logits.shape

        prob = out_logits.sigmoid()
        scores = prob.reshape(bs, -1)

        # convert to [x0, y0, x1, y1] format
        boxes = box_cxcywh_to_xyxy(out_bbox)
        boxes = boxes.unsqueeze(2).repeat(1, 1, category, 1).reshape(bs, -1, 4)
        # and from relative [0, 1] to absolute [0, height] coordinates
        img_h, img_w = target_sizes.unbind(1)
        scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1)
        boxes = boxes * scale_fct[:, None, :]

        results = []
        for i in range(bs):
            result = dict()
            scores_per_image, topk_indices = scores[i].topk(query_len, sorted=False)
            labels_per_image = topk_indices % category
            boxes_per_image = boxes[i][topk_indices]
            result['scores'] = scores_per_image
            result['labels'] = labels_per_image
            result['boxes'] = boxes_per_image
            results.append(result)

def bias_init_with_prob(prior_prob):
    """initialize conv/fc bias value according to giving probablity."""
    bias_init = float(-np.log((1 - prior_prob) / prior_prob))
    return bias_init